#include <stdio.h>
#include <unistd.h>
#include <string.h>

#include "security.h"
#include "security_error.h"
#include "security_common.h"

#include "mbedtls/des.h"

#define PKCS7_BLOCK_SIZE 8

static int internal_des_encrypt_file(FILE *fp_in, FILE *fp_out, char key[DES_KEY_SIZE]) {
    char buffer_in[PKCS7_BLOCK_SIZE], buffer_out[PKCS7_BLOCK_SIZE];
    int8_t buf_in_len = 0, buf_out_len = 0, write_len = 0;
    int8_t i = 0;
    int ret = SECURITY_OK;

    bool exit_flag = false;

    mbedtls_des_context ctx;

    mbedtls_des_init(&ctx);

    ret = mbedtls_des_setkey_enc(&ctx, key);
    if(ret != 0) {
        ret = SECURITY_DES_ENCRYPT_ERROR;
        goto exit;
    }

    while(1) {
        memset(buffer_in, 0, PKCS7_BLOCK_SIZE);
        memset(buffer_out, 0, PKCS7_BLOCK_SIZE);

        // 数据从文件中读取出来
        if(feof(fp_in) != 0) {  // 如果读到末尾，而且没有退出，则证明上一次读取正好读完八个字节，所以需要填充八位进行标记
            memset(buffer_in, 0x08, PKCS7_BLOCK_SIZE);  // 进行PKCS#7填充
            buf_in_len = PKCS7_BLOCK_SIZE;
            exit_flag = true;
        }
        else {  // 一般情况
            buf_in_len = fread(buffer_in, sizeof(char), PKCS7_BLOCK_SIZE, fp_in);
            if(buf_in_len != PKCS7_BLOCK_SIZE) {    // 如果读不满则代表读完了
                for(i = buf_in_len; i < PKCS7_BLOCK_SIZE; i++) {    // 进行PKCS#7填充
                    buffer_in[i] = PKCS7_BLOCK_SIZE - buf_in_len;
                    exit_flag = true;
                }
            }
            buf_in_len = PKCS7_BLOCK_SIZE;
        }

        // 加密
        ret = mbedtls_des_crypt_ecb(&ctx, buffer_in, buffer_out);
        if(ret != 0) {
            ret = SECURITY_DES_ENCRYPT_ERROR;
            goto exit;
        }

        // 保存到文件中
        write_len = fwrite(buffer_out, sizeof(char), PKCS7_BLOCK_SIZE, fp_out);
        if(PKCS7_BLOCK_SIZE != write_len) {
            ret = SECURITY_ERR_FILE_WRITE_ERROR;
            goto exit;
        }

        // 询问标志位是否退出循环
        if(true == exit_flag) {
            break;
        }
    }

    ret = SECURITY_OK;

exit:
    mbedtls_des_free(&ctx);

    return ret;
}

static int internal_des_decrypt_file(FILE *fp_in, FILE *fp_out, char key[DES_KEY_SIZE]) {
    char buffer_in[PKCS7_BLOCK_SIZE], buffer_out[PKCS7_BLOCK_SIZE];
    int8_t buf_in_len = 0, buf_out_len = 0, write_len = 0;
    int8_t i = 0;
    int ret = SECURITY_OK;
    int file_size = 0;

    mbedtls_des_context ctx;

    mbedtls_des_init(&ctx);

    ret = mbedtls_des_setkey_dec(&ctx, key);
    if(ret != 0) {
        ret = SECURITY_DES_ENCRYPT_ERROR;
        goto exit;
    }

    // 拿到文件大小
    fseek(fp_in, 0, SEEK_END);
    file_size = ftell(fp_in);

    // 重新回到开始位置
    fseek(fp_in, 0, SEEK_SET);

    while(1) {
        memset(buffer_in, 0, PKCS7_BLOCK_SIZE);
        memset(buffer_out, 0, PKCS7_BLOCK_SIZE);

        // 从文件中读出数据，由于在之前的加密时保证了绝对的八字节一组，所以读出的数据长度必定是八字节一组
        buf_in_len = fread(buffer_in, sizeof(char), PKCS7_BLOCK_SIZE, fp_in);
        if(0 == buf_in_len) {   // 如果读出0字节内容则代表已经完全处理完了
            ret = SECURITY_OK;
            goto exit;
        }

        // 解密
        mbedtls_des_crypt_ecb(&ctx, buffer_in, buffer_out);
        if(ret != 0) {
            ret = SECURITY_DES_DECRYPT_ERROR;
            goto exit;
        }

        // 如果读到了末尾，则需要解填充
        if(feof(fp_in) != 0 || ftell(fp_in) == file_size) {
            if((unsigned char)buffer_out[PKCS7_BLOCK_SIZE - 1] > PKCS7_BLOCK_SIZE) {   // 当最后一位填充大于块大小的时候视为解密错误
                ret = SECURITY_DES_DECRYPT_ERROR;
                goto exit;
            }
            for(i = PKCS7_BLOCK_SIZE - 1; i >= PKCS7_BLOCK_SIZE - buffer_out[PKCS7_BLOCK_SIZE - 1]; i--) {  // 当不符合PKCS#7填充规则时也视为解密错误
                if(buffer_out[PKCS7_BLOCK_SIZE - 1] != buffer_out[i]) {
                    ret = SECURITY_DES_DECRYPT_ERROR;
                    goto exit;
                }
            }

            // 解密成功的情况
            if(buffer_out[PKCS7_BLOCK_SIZE - 1] != PKCS7_BLOCK_SIZE) {
                // 写入到文件中
                write_len = fwrite(buffer_out, sizeof(char), PKCS7_BLOCK_SIZE - buffer_out[PKCS7_BLOCK_SIZE - 1], fp_out);
                if((PKCS7_BLOCK_SIZE - buffer_out[PKCS7_BLOCK_SIZE - 1]) != write_len) {
                    ret = SECURITY_ERR_FILE_WRITE_ERROR;
                    goto exit;
                }
            }
            ret = SECURITY_OK;
            goto exit;
        }

        write_len = fwrite(buffer_out, sizeof(char), PKCS7_BLOCK_SIZE, fp_out);
        if(PKCS7_BLOCK_SIZE != write_len) {
            ret = SECURITY_ERR_FILE_WRITE_ERROR;
            goto exit;
        }

    }

    ret = SECURITY_OK;

exit:
    mbedtls_des_free(&ctx);

    return ret;
}


int security_des_encrypt(char out[], size_t *olen, size_t osize, char in[], size_t inlen, char key[DES_KEY_SIZE]) {
    mbedtls_des_context ctx;
    int ret = 0, count = 0;
    char *padding = NULL;
    size_t p_len = 0;
    size_t p_size = 0;

    if(out == NULL || olen == NULL || in == NULL || key == NULL || osize == 0 || inlen == 0) {
        return SECURITY_ERR_BAD_INPUT;
    }

    if(osize < inlen + 9) {
        return SECURITY_ERR_BUFFER_TOO_SMALL;
    }

    padding = (char *)malloc(sizeof(char) * (inlen + 9));
    p_size = inlen + 9;
    if(padding == NULL) {
        ret = SECURITY_ERR_MALLOC;
        goto exit;
    }
    ret = padding_pkcs7(padding, &p_len, p_size, in, inlen, PKCS7_BLOCK_SIZE);
    if(ret != INTERNAL_SECURITY_OK) {
        ret = SECURITY_ERR_PADDING;
        goto exit;
    }

    mbedtls_des_init(&ctx);
    ret = mbedtls_des_setkey_enc(&ctx, key);
    if(ret != 0) {
        ret = SECURITY_DES_ENCRYPT_ERROR;
        goto exit;
    }

    count = 0;
    while(count < p_len) {
        ret = mbedtls_des_crypt_ecb(&ctx, padding + count, out + count);
        if(ret != 0) {
            ret = SECURITY_DES_ENCRYPT_ERROR;
            goto exit;
        }
        count += 8;
    }

    *olen = count;

exit:
    mbedtls_des_free(&ctx);
    if(padding != NULL) {
        free(padding);
    }
    return ret;
}

int security_des_decrypt(char out[], size_t *olen, size_t osize, char in[], size_t inlen, char key[DES_KEY_SIZE]) {
    mbedtls_des_context ctx;
    int ret = 0, count = 0;

    if(out == NULL || olen == NULL || in == NULL || key == NULL || osize == 0 || inlen == 0) {
        return SECURITY_ERR_BAD_INPUT;
    }

    if(inlen % 8 != 0) {
        return SECURITY_ERR_BAD_INPUT;
    }

    if(osize < inlen) {
        return SECURITY_ERR_BUFFER_TOO_SMALL;
    }

    mbedtls_des_init(&ctx);

    ret = mbedtls_des_setkey_dec(&ctx, key);
    if(ret != 0) {
        ret = SECURITY_DES_ENCRYPT_ERROR;
        goto exit;
    }

    while(count < inlen) {
        ret = mbedtls_des_crypt_ecb(&ctx, in + count, out + count);
        if(ret != 0) {
            ret = SECURITY_DES_DECRYPT_ERROR;
            goto exit;
        }
        count += 8;
    }

    *olen = count;

    ret = unpadding_pkcs7(out, olen, osize, out, count);
    if(ret != INTERNAL_SECURITY_OK) {
        ret = SECURITY_ERR_PADDING;
        goto exit;
    }

exit:
    mbedtls_des_free(&ctx);
    return ret;
}

int security_des_file_encrypt(const char *input_file, const char *output_file, char key[DES_KEY_SIZE]) {

    FILE *fp_input = NULL, *fp_output = NULL;

    int ret = SECURITY_OK;

    if(input_file == NULL || output_file == NULL || key == NULL) {
        ret = SECURITY_ERR_BAD_INPUT;
        goto exit;
    }

    // 检测文件是否有读取权限
    if(access(input_file, R_OK) != 0) {
        ret = SECURITY_ERR_FILE_CANNOT_READ;
        goto exit;
    }

    // 打开input文件
    fp_input = fopen(input_file, "rb");
    if(NULL == fp_input) {
        ret = SECURITY_ERR_FILE_OPEN;
        goto exit;
    }

    // 打开output文件
    fp_output = fopen(output_file, "wb");
    if(NULL == fp_output) {
        ret = SECURITY_ERR_FILE_OPEN;
        goto exit;
    }

    ret = internal_des_encrypt_file(fp_input, fp_output, key);

exit:
    if(fp_input != NULL) {
        fclose(fp_input);
    }

    if(fp_output != NULL) {
        fclose(fp_output);
    }

    return ret;
}

int security_des_file_decrypt(const char *input_file, const char *output_file, char key[DES_KEY_SIZE]) {
    FILE *fp_input = NULL, *fp_output = NULL;
    int ret = SECURITY_OK;

    if(input_file == NULL || output_file == NULL || key == NULL) {
        ret = SECURITY_ERR_BAD_INPUT;
        goto exit;
    }

        // 检测文件是否有读取权限
    if(access(input_file, R_OK) != 0) {
        ret = SECURITY_ERR_FILE_CANNOT_READ;
        goto exit;
    }

    // 打开input文件
    fp_input = fopen(input_file, "rb");
    if(NULL == fp_input) {
        ret = SECURITY_ERR_FILE_OPEN;
        goto exit;
    }

    // 打开output文件
    fp_output = fopen(output_file, "wb");
    if(NULL == fp_output) {
        ret = SECURITY_ERR_FILE_OPEN;
        goto exit;
    }

    ret = internal_des_decrypt_file(fp_input, fp_output, key);

exit:
    if(fp_input != NULL) {
        fclose(fp_input);
    }

    if(fp_output != NULL) {
        fclose(fp_output);
    }

    return ret;
    
}